import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.inspection import permutation_importance
import matplotlib.pyplot as plt
import seaborn as sns


# 解决中文乱码问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 数据加载与预处理
data = pd.read_csv("../data/train.csv")
data.drop(columns=['EmployeeNumber', 'StandardHours', 'Over18'], inplace=True)

# 处理分类特征
categorical_cols = data.select_dtypes(include=['object']).columns
for col in categorical_cols:
    le = LabelEncoder()
    data[col] = le.fit_transform(data[col])  # 将分类特征转换为数值型

# 分离特征和目标变量
X = data.drop('Attrition', axis=1)
y = data['Attrition']

# 训练随机森林
rf = RandomForestClassifier(n_estimators=100, random_state=42, oob_score=True)
rf.fit(X, y)

# 基尼重要性
gini_importance = pd.DataFrame({
    'Feature': X.columns,
    'Gini_Importance': rf.feature_importances_
}).sort_values('Gini_Importance', ascending=False)

# 置换重要性
perm_importance = permutation_importance(
    rf, X, y, n_repeats=10, random_state=42
)
perm_df = pd.DataFrame({
    'Feature': X.columns,
    'Perm_Importance': perm_importance.importances_mean
}).sort_values('Perm_Importance', ascending=False)

# 合并结果
result_df = pd.merge(gini_importance, perm_df, on='Feature')

# 关键优化：交换坐标轴的水平条形图
plt.figure(figsize=(12, 10))  # 增大高度以适应更多特征名

# 创建水平条形图（交换X/Y轴）
ax = sns.barplot(
    x='Importance',  # 重要性得分在X轴
    y='Feature',     # 特征名在Y轴
    data=result_df.melt(id_vars='Feature',
                        value_vars=['Gini_Importance', 'Perm_Importance'],
                        var_name='Importance_Type',
                        value_name='Importance'),
    hue='Importance_Type',  # 按重要性类型着色
    palette=['#1f77b4', '#ff7f0e'],  # 定制颜色
    orient='h'  # 水平方向[6,8](@ref)
)

# 添加数值标签（提高可读性）
for p in ax.patches:
    width = p.get_width()
    plt.text(width + 0.005,
             p.get_y() + p.get_height()/2,
             f'{width:.3f}',
             ha='left',
             va='center',
             fontsize=9)

# 图表美化
plt.title("特征重要性对比分析", fontsize=16)
plt.xlabel("重要性得分", fontsize=12)
plt.ylabel("")  # 隐藏Y轴标题
plt.xticks(rotation=0)  # X轴标签保持水平
plt.grid(axis='x', linestyle='--', alpha=0.7)  # 添加辅助网格线
plt.legend(title='重要性类型',
           labels=['基尼重要性', '置换重要性'],
           loc='lower right')  # 图例位置
sns.despine(left=True)  # 隐藏左侧边框

plt.tight_layout()
# plt.savefig("../data/fig/feature_importance.png")
plt.show()